{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### BabyAGI\n",
    "\n",
    "An implementation of BabyAGI using LangChain and OpenAI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: langchain in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (0.0.138)\n",
      "Requirement already satisfied: SQLAlchemy<2,>=1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (1.4.47)\n",
      "Requirement already satisfied: pydantic<2,>=1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (1.10.6)\n",
      "Requirement already satisfied: PyYAML>=5.4.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (6.0)\n",
      "Requirement already satisfied: openapi-schema-pydantic<2.0,>=1.2 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (1.2.4)\n",
      "Requirement already satisfied: dataclasses-json<0.6.0,>=0.5.7 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (0.5.7)\n",
      "Requirement already satisfied: numpy<2,>=1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (1.23.5)\n",
      "Requirement already satisfied: requests<3,>=2 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (2.28.2)\n",
      "Requirement already satisfied: tenacity<9.0.0,>=8.1.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (8.2.2)\n",
      "Requirement already satisfied: aiohttp<4.0.0,>=3.8.3 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (3.8.4)\n",
      "Requirement already satisfied: async-timeout<5.0.0,>=4.0.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (4.0.2)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (1.3.3)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (1.8.2)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (22.2.0)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (1.3.1)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (6.0.4)\n",
      "Requirement already satisfied: charset-normalizer<4.0,>=2.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (3.1.0)\n",
      "Requirement already satisfied: marshmallow-enum<2.0.0,>=1.5.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from dataclasses-json<0.6.0,>=0.5.7->langchain) (1.5.1)\n",
      "Requirement already satisfied: typing-inspect>=0.4.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from dataclasses-json<0.6.0,>=0.5.7->langchain) (0.8.0)\n",
      "Requirement already satisfied: marshmallow<4.0.0,>=3.3.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from dataclasses-json<0.6.0,>=0.5.7->langchain) (3.19.0)\n",
      "Requirement already satisfied: typing-extensions>=4.2.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from pydantic<2,>=1->langchain) (4.5.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from requests<3,>=2->langchain) (3.4)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from requests<3,>=2->langchain) (1.26.15)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from requests<3,>=2->langchain) (2022.12.7)\n",
      "Requirement already satisfied: greenlet!=0.4.17 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from SQLAlchemy<2,>=1->langchain) (2.0.2)\n",
      "Requirement already satisfied: packaging>=17.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from marshmallow<4.0.0,>=3.3.0->dataclasses-json<0.6.0,>=0.5.7->langchain) (23.0)\n",
      "Requirement already satisfied: mypy-extensions>=0.3.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from typing-inspect>=0.4.0->dataclasses-json<0.6.0,>=0.5.7->langchain) (1.0.0)\n",
      "Requirement already satisfied: openai in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (0.27.4)\n",
      "Requirement already satisfied: aiohttp in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from openai) (3.8.4)\n",
      "Requirement already satisfied: requests>=2.20 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from openai) (2.28.2)\n",
      "Requirement already satisfied: tqdm in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from openai) (4.65.0)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from requests>=2.20->openai) (3.1.0)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from requests>=2.20->openai) (1.26.15)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from requests>=2.20->openai) (3.4)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from requests>=2.20->openai) (2022.12.7)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from aiohttp->openai) (1.3.1)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from aiohttp->openai) (1.3.3)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from aiohttp->openai) (22.2.0)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from aiohttp->openai) (4.0.2)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from aiohttp->openai) (1.8.2)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from aiohttp->openai) (6.0.4)\n",
      "Requirement already satisfied: unstructured in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (0.5.12)\n",
      "Requirement already satisfied: lxml in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from unstructured) (4.9.2)\n",
      "Requirement already satisfied: pandas in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from unstructured) (1.5.3)\n",
      "Requirement already satisfied: openpyxl in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from unstructured) (3.1.2)\n",
      "Requirement already satisfied: pillow in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from unstructured) (9.4.0)\n",
      "Requirement already satisfied: nltk in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from unstructured) (3.8.1)\n",
      "Requirement already satisfied: markdown in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from unstructured) (3.4.3)\n",
      "Requirement already satisfied: python-magic in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from unstructured) (0.4.27)\n",
      "Requirement already satisfied: python-docx in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from unstructured) (0.8.11)\n",
      "Requirement already satisfied: python-pptx in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from unstructured) (0.6.21)\n",
      "Requirement already satisfied: requests in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from unstructured) (2.28.2)\n",
      "Requirement already satisfied: msg-parser in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from unstructured) (1.2.0)\n",
      "Requirement already satisfied: certifi>=2022.12.07 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from unstructured) (2022.12.7)\n",
      "Requirement already satisfied: argilla in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from unstructured) (1.5.1)\n",
      "Requirement already satisfied: pypandoc in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from unstructured) (1.11)\n",
      "Requirement already satisfied: wrapt<1.15,>=1.13 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from argilla->unstructured) (1.14.1)\n",
      "Requirement already satisfied: tqdm>=4.27.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from argilla->unstructured) (4.65.0)\n",
      "Requirement already satisfied: rich<=13.0.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from argilla->unstructured) (13.0.1)\n",
      "Requirement already satisfied: backoff in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from argilla->unstructured) (2.2.1)\n",
      "Requirement already satisfied: packaging>=20.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from argilla->unstructured) (23.0)\n",
      "Requirement already satisfied: httpx<0.24,>=0.15 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from argilla->unstructured) (0.23.3)\n",
      "Requirement already satisfied: pydantic>=1.7.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from argilla->unstructured) (1.10.6)\n",
      "Requirement already satisfied: monotonic in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from argilla->unstructured) (1.6)\n",
      "Requirement already satisfied: numpy<1.24.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from argilla->unstructured) (1.23.5)\n",
      "Requirement already satisfied: deprecated~=1.2.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from argilla->unstructured) (1.2.13)\n",
      "Requirement already satisfied: python-dateutil>=2.8.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from pandas->unstructured) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2020.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from pandas->unstructured) (2022.7.1)\n",
      "Requirement already satisfied: olefile>=0.46 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from msg-parser->unstructured) (0.46)\n",
      "Requirement already satisfied: click in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from nltk->unstructured) (8.1.3)\n",
      "Requirement already satisfied: joblib in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from nltk->unstructured) (1.2.0)\n",
      "Requirement already satisfied: regex>=2021.8.3 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from nltk->unstructured) (2022.10.31)\n",
      "Requirement already satisfied: et-xmlfile in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from openpyxl->unstructured) (1.1.0)\n",
      "Requirement already satisfied: XlsxWriter>=0.5.7 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from python-pptx->unstructured) (3.0.9)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from requests->unstructured) (3.4)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from requests->unstructured) (1.26.15)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from requests->unstructured) (3.1.0)\n",
      "Requirement already satisfied: rfc3986[idna2008]<2,>=1.3 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from httpx<0.24,>=0.15->argilla->unstructured) (1.5.0)\n",
      "Requirement already satisfied: httpcore<0.17.0,>=0.15.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from httpx<0.24,>=0.15->argilla->unstructured) (0.16.3)\n",
      "Requirement already satisfied: sniffio in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from httpx<0.24,>=0.15->argilla->unstructured) (1.3.0)\n",
      "Requirement already satisfied: typing-extensions>=4.2.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from pydantic>=1.7.1->argilla->unstructured) (4.5.0)\n",
      "Requirement already satisfied: six>=1.5 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from python-dateutil>=2.8.1->pandas->unstructured) (1.16.0)\n",
      "Requirement already satisfied: commonmark<0.10.0,>=0.9.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from rich<=13.0.1->argilla->unstructured) (0.9.1)\n",
      "Requirement already satisfied: pygments<3.0.0,>=2.6.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from rich<=13.0.1->argilla->unstructured) (2.14.0)\n",
      "Requirement already satisfied: h11<0.15,>=0.13 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from httpcore<0.17.0,>=0.15.0->httpx<0.24,>=0.15->argilla->unstructured) (0.14.0)\n",
      "Requirement already satisfied: anyio<5.0,>=3.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from httpcore<0.17.0,>=0.15.0->httpx<0.24,>=0.15->argilla->unstructured) (3.6.2)\n",
      "Requirement already satisfied: pypdf in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (3.7.1)\n",
      "Requirement already satisfied: faiss-gpu in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (1.7.2)\n"
     ]
    }
   ],
   "source": [
    "! pip install langchain --upgrade\n",
    "! pip install openai --upgrade\n",
    "! pip install unstructured --upgrade\n",
    "! pip install pypdf --upgrade\n",
    "! pip install faiss-gpu --upgrade"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import Packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import faiss\n",
    "import textwrap\n",
    "\n",
    "from collections import deque\n",
    "from typing import Dict, List, Optional, Any\n",
    "\n",
    "from langchain import LLMChain, OpenAI, PromptTemplate\n",
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "from langchain.llms import BaseLLM\n",
    "from langchain.vectorstores.base import VectorStore\n",
    "from pydantic import BaseModel, Field\n",
    "from langchain.chains.base import Chain\n",
    "from langchain.vectorstores import FAISS\n",
    "from langchain.docstore import InMemoryDocstore"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### OpenAI API Key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = input(\"OpenAI API Key: \")\n",
    "os.environ[\"OPENAI_API_KEY\"] = query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings_model = OpenAIEmbeddings()\n",
    "embedding_size = 1536\n",
    "index = faiss.IndexFlatL2(embedding_size)\n",
    "vectorstore = FAISS(embeddings_model.embed_query, index, InMemoryDocstore({}), {})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TaskCreationChain(LLMChain):\n",
    "    \"\"\"Chain to generates tasks.\"\"\"\n",
    "\n",
    "    @classmethod\n",
    "    def from_llm(cls, llm: BaseLLM, verbose: bool = True) -> LLMChain:\n",
    "        \"\"\"Get the response parser.\"\"\"\n",
    "        task_creation_template = (\n",
    "            \"You are an task creation AI that uses the result of an execution agent\"\n",
    "            \" to create new tasks with the following objective: {objective},\"\n",
    "            \" The last completed task has the result: {result}.\"\n",
    "            \" This result was based on this task description: {task_description}.\"\n",
    "            \" These are incomplete tasks: {incomplete_tasks}.\"\n",
    "            \" Based on the result, create new tasks to be completed\"\n",
    "            \" by the AI system that do not overlap with incomplete tasks.\"\n",
    "            \" Return the tasks as an array.\"\n",
    "        )\n",
    "        prompt = PromptTemplate(\n",
    "            template=task_creation_template,\n",
    "            input_variables=[\"result\", \"task_description\", \"incomplete_tasks\", \"objective\"],\n",
    "        )\n",
    "        return cls(prompt=prompt, llm=llm, verbose=verbose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TaskPrioritizationChain(LLMChain):\n",
    "    \"\"\"Chain to prioritize tasks.\"\"\"\n",
    "\n",
    "    @classmethod\n",
    "    def from_llm(cls, llm: BaseLLM, verbose: bool = True) -> LLMChain:\n",
    "        \"\"\"Get the response parser.\"\"\"\n",
    "        task_prioritization_template = (\n",
    "            \"You are an task prioritization AI tasked with cleaning the formatting of and reprioritizing\"\n",
    "            \" the following tasks: {task_names}.\"\n",
    "            \" Consider the ultimate objective of your team: {objective}.\"\n",
    "            \" Do not remove any tasks. Return the result as a numbered list, like:\"\n",
    "            \" #. First task\"\n",
    "            \" #. Second task\"\n",
    "            \" Start the task list with number {next_task_id}.\"\n",
    "        )\n",
    "        prompt = PromptTemplate(\n",
    "            template=task_prioritization_template,\n",
    "            input_variables=[\"task_names\", \"next_task_id\", \"objective\"],\n",
    "        )\n",
    "        return cls(prompt=prompt, llm=llm, verbose=verbose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ExecutionChain(LLMChain):\n",
    "    \"\"\"Chain to execute tasks.\"\"\"\n",
    "\n",
    "    @classmethod\n",
    "    def from_llm(cls, llm: BaseLLM, verbose: bool = True) -> LLMChain:\n",
    "        \"\"\"Get the response parser.\"\"\"\n",
    "        execution_template = (\n",
    "            \"You are an AI who performs one task based on the following objective: {objective}.\"\n",
    "            \" Take into account these previously completed tasks: {context}.\"\n",
    "            \" Your task: {task}.\"\n",
    "            \" Response:\"\n",
    "        )\n",
    "        prompt = PromptTemplate(\n",
    "            template=execution_template,\n",
    "            input_variables=[\"objective\", \"context\", \"task\"],\n",
    "        )\n",
    "        return cls(prompt=prompt, llm=llm, verbose=verbose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_next_task(task_creation_chain: LLMChain, result: Dict, task_description: str, task_list: List[str], objective: str) -> List[Dict]:\n",
    "    \"\"\"Get the next task.\"\"\"\n",
    "    incomplete_tasks = \", \".join(task_list)\n",
    "    response = task_creation_chain.run(result=result, task_description=task_description, incomplete_tasks=incomplete_tasks, objective=objective)\n",
    "    new_tasks = response.split('\\n')\n",
    "    return [{\"task_name\": task_name} for task_name in new_tasks if task_name.strip()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prioritize_tasks(task_prioritization_chain: LLMChain, this_task_id: int, task_list: List[Dict], objective: str) -> List[Dict]:\n",
    "    \"\"\"Prioritize tasks.\"\"\"\n",
    "    task_names = [t[\"task_name\"] for t in task_list]\n",
    "    next_task_id = int(this_task_id) + 1\n",
    "    response = task_prioritization_chain.run(task_names=task_names, next_task_id=next_task_id, objective=objective)\n",
    "    new_tasks = response.split('\\n')\n",
    "    prioritized_task_list = []\n",
    "    for task_string in new_tasks:\n",
    "        if not task_string.strip():\n",
    "            continue\n",
    "        task_parts = task_string.strip().split(\".\", 1)\n",
    "        if len(task_parts) == 2:\n",
    "            task_id = task_parts[0].strip()\n",
    "            task_name = task_parts[1].strip()\n",
    "            prioritized_task_list.append({\"task_id\": task_id, \"task_name\": task_name})\n",
    "    return prioritized_task_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _get_top_tasks(vectorstore, query: str, k: int) -> List[str]:\n",
    "    \"\"\"Get the top k tasks based on the query.\"\"\"\n",
    "    results = vectorstore.similarity_search_with_score(query, k=k)\n",
    "    if not results:\n",
    "        return []\n",
    "    sorted_results, _ = zip(*sorted(results, key=lambda x: x[1], reverse=True))\n",
    "    return [str(item.metadata['task']) for item in sorted_results]\n",
    "\n",
    "def execute_task(vectorstore, execution_chain: LLMChain, objective: str, task: str, k: int = 5) -> str:\n",
    "    \"\"\"Execute a task.\"\"\"\n",
    "    context = _get_top_tasks(vectorstore, query=objective, k=k)\n",
    "    return execution_chain.run(objective=objective, context=context, task=task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BabyAGI(Chain, BaseModel):\n",
    "    \"\"\"Controller model for the BabyAGI agent.\"\"\"\n",
    "\n",
    "    task_list: deque = Field(default_factory=deque)\n",
    "    task_creation_chain: TaskCreationChain = Field(...)\n",
    "    task_prioritization_chain: TaskPrioritizationChain = Field(...)\n",
    "    execution_chain: ExecutionChain = Field(...)\n",
    "    task_id_counter: int = Field(1)\n",
    "    vectorstore: VectorStore = Field(init=False)\n",
    "    max_iterations: Optional[int] = None\n",
    "        \n",
    "    class Config:\n",
    "        \"\"\"Configuration for this pydantic object.\"\"\"\n",
    "        arbitrary_types_allowed = True\n",
    "\n",
    "    def add_task(self, task: Dict):\n",
    "        self.task_list.append(task)\n",
    "\n",
    "    def print_task_list(self):\n",
    "        print(\"\\033[95m\\033[1m\" + \"\\n*****TASK LIST*****\\n\" + \"\\033[0m\\033[0m\")\n",
    "        for t in self.task_list:\n",
    "            print(str(t[\"task_id\"]) + \": \" + t[\"task_name\"])\n",
    "\n",
    "    def print_next_task(self, task: Dict):\n",
    "        print(\"\\033[92m\\033[1m\" + \"\\n*****NEXT TASK*****\\n\" + \"\\033[0m\\033[0m\")\n",
    "        print(str(task[\"task_id\"]) + \": \" + task[\"task_name\"])\n",
    "\n",
    "    def print_task_result(self, result: str):\n",
    "        print(\"\\033[93m\\033[1m\" + \"\\n*****TASK RESULT*****\\n\" + \"\\033[0m\\033[0m\")\n",
    "        print(result)\n",
    "        \n",
    "    @property\n",
    "    def input_keys(self) -> List[str]:\n",
    "        return [\"objective\"]\n",
    "    \n",
    "    @property\n",
    "    def output_keys(self) -> List[str]:\n",
    "        return []\n",
    "\n",
    "    def _call(self, inputs: Dict[str, Any]) -> Dict[str, Any]:\n",
    "        \"\"\"Run the agent.\"\"\"\n",
    "        objective = inputs['objective']\n",
    "        first_task = inputs.get(\"first_task\", \"Make a todo list\")\n",
    "        self.add_task({\"task_id\": 1, \"task_name\": first_task})\n",
    "        num_iters = 0\n",
    "        while True:\n",
    "            if self.task_list:\n",
    "                self.print_task_list()\n",
    "\n",
    "                # Step 1: Pull the first task\n",
    "                task = self.task_list.popleft()\n",
    "                self.print_next_task(task)\n",
    "\n",
    "                # Step 2: Execute the task\n",
    "                result = execute_task(\n",
    "                    self.vectorstore, self.execution_chain, objective, task[\"task_name\"]\n",
    "                )\n",
    "                this_task_id = int(task[\"task_id\"])\n",
    "                self.print_task_result(result)\n",
    "\n",
    "                # Step 3: Store the result in Pinecone\n",
    "                result_id = f\"result_{task['task_id']}\"\n",
    "                self.vectorstore.add_texts(\n",
    "                    texts=[result],\n",
    "                    metadatas=[{\"task\": task[\"task_name\"]}],\n",
    "                    ids=[result_id],\n",
    "                )\n",
    "\n",
    "                # Step 4: Create new tasks and reprioritize task list\n",
    "                new_tasks = get_next_task(\n",
    "                    self.task_creation_chain, result, task[\"task_name\"], [t[\"task_name\"] for t in self.task_list], objective\n",
    "                )\n",
    "                for new_task in new_tasks:\n",
    "                    self.task_id_counter += 1\n",
    "                    new_task.update({\"task_id\": self.task_id_counter})\n",
    "                    self.add_task(new_task)\n",
    "                self.task_list = deque(\n",
    "                    prioritize_tasks(\n",
    "                        self.task_prioritization_chain, this_task_id, list(self.task_list), objective\n",
    "                    )\n",
    "                )\n",
    "            num_iters += 1\n",
    "            if self.max_iterations is not None and num_iters == self.max_iterations:\n",
    "                print(\"\\033[91m\\033[1m\" + \"\\n*****TASK ENDING*****\\n\" + \"\\033[0m\\033[0m\")\n",
    "                break\n",
    "        return {}\n",
    "\n",
    "    @classmethod\n",
    "    def from_llm(\n",
    "        cls,\n",
    "        llm: BaseLLM,\n",
    "        vectorstore: VectorStore,\n",
    "        verbose: bool = False,\n",
    "        **kwargs\n",
    "    ) -> \"BabyAGI\":\n",
    "        \"\"\"Initialize the BabyAGI Controller.\"\"\"\n",
    "        task_creation_chain = TaskCreationChain.from_llm(\n",
    "            llm, verbose=verbose\n",
    "        )\n",
    "        task_prioritization_chain = TaskPrioritizationChain.from_llm(\n",
    "            llm, verbose=verbose\n",
    "        )\n",
    "        execution_chain = ExecutionChain.from_llm(llm, verbose=verbose)\n",
    "        return cls(\n",
    "            task_creation_chain=task_creation_chain,\n",
    "            task_prioritization_chain=task_prioritization_chain,\n",
    "            execution_chain=execution_chain,\n",
    "            vectorstore=vectorstore,\n",
    "            **kwargs\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(temperature=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Logging of LLMChains\n",
    "verbose=False\n",
    "# If None, will keep on going forever\n",
    "max_iterations: Optional[int] = 3\n",
    "baby_agi = BabyAGI.from_llm(\n",
    "    llm=llm,\n",
    "    vectorstore=vectorstore,\n",
    "    verbose=verbose,\n",
    "    max_iterations=max_iterations\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Human:  Plan a travel itinerary to boracay for 5 days. Must include outdoor\n",
      "activities and eating filipino foods.\n",
      "\u001b[95m\u001b[1m\n",
      "*****TASK LIST*****\n",
      "\u001b[0m\u001b[0m\n",
      "1: Make a todo list\n",
      "\u001b[92m\u001b[1m\n",
      "*****NEXT TASK*****\n",
      "\u001b[0m\u001b[0m\n",
      "1: Make a todo list\n",
      "\u001b[93m\u001b[1m\n",
      "*****TASK RESULT*****\n",
      "\u001b[0m\u001b[0m\n",
      "\n",
      "\n",
      "1. Book flights to Boracay\n",
      "2. Book accommodation in Boracay\n",
      "3. Research outdoor activities in Boracay\n",
      "4. Make a list of Filipino restaurants to visit\n",
      "5. Make a list of must-see attractions in Boracay\n",
      "6. Create a budget for the trip\n",
      "7. Pack necessary items for the trip\n",
      "8. Arrange transportation to and from the airport\n",
      "9. Arrange transportation to and from outdoor activities\n",
      "10. Arrange transportation to and from restaurants and attractions\n",
      "\u001b[95m\u001b[1m\n",
      "*****TASK LIST*****\n",
      "\u001b[0m\u001b[0m\n",
      "2: Create a timeline for the trip\n",
      "3: Create a map of Boracay\n",
      "4: Research the cost of outdoor activities in Boracay\n",
      "5: Research the cost of transportation to and from outdoor activities\n",
      "6: Research the cost of transportation to and from restaurants and attractions\n",
      "7: Research the cost of food at Filipino restaurants\n",
      "8: Research the cost of attractions in Boracay\n",
      "9: Create a list of items to bring on the trip\n",
      "10: Create a list of emergency contacts\n",
      "11: Create a list of emergency items to bring on the trip\n",
      "\u001b[92m\u001b[1m\n",
      "*****NEXT TASK*****\n",
      "\u001b[0m\u001b[0m\n",
      "2: Create a timeline for the trip\n",
      "\u001b[93m\u001b[1m\n",
      "*****TASK RESULT*****\n",
      "\u001b[0m\u001b[0m\n",
      "\n",
      "\n",
      "Day 1:\n",
      "- Arrive in Boracay\n",
      "- Check into hotel\n",
      "- Explore White Beach\n",
      "- Enjoy dinner at a local Filipino restaurant\n",
      "\n",
      "Day 2:\n",
      "- Take a boat tour around the island\n",
      "- Visit Puka Beach\n",
      "- Enjoy lunch at a local Filipino restaurant\n",
      "\n",
      "Day 3:\n",
      "- Take a parasailing tour\n",
      "- Visit Bulabog Beach\n",
      "- Enjoy dinner at a local Filipino restaurant\n",
      "\n",
      "Day 4:\n",
      "- Take a scuba diving tour\n",
      "- Visit Ilig-Iligan Beach\n",
      "- Enjoy lunch at a local Filipino restaurant\n",
      "\n",
      "Day 5:\n",
      "- Take a windsurfing tour\n",
      "- Visit Diniwid Beach\n",
      "- Enjoy dinner at a local Filipino restaurant\n",
      "- Depart from Boracay\n",
      "\u001b[95m\u001b[1m\n",
      "*****TASK LIST*****\n",
      "\u001b[0m\u001b[0m\n",
      "3: Create a map of Boracay\n",
      "4: Research the cost of outdoor activities in Boracay\n",
      "5: Research the cost of transportation to and from outdoor activities\n",
      "6: Research the cost of transportation to and from restaurants and attractions\n",
      "7: Research the cost of food at Filipino restaurants\n",
      "8: Research the cost of attractions in Boracay\n",
      "9: Create a list of items to bring on the trip\n",
      "10: Create a list of emergency contacts\n",
      "11: Create a list of emergency items to bring on the trip\n",
      "12: Research the best time of year to visit Boracay\n",
      "13: Research the best time of day to visit outdoor activities\n",
      "14: Research the best time of day to visit attractions\n",
      "15: Research the best time of day to visit restaurants\n",
      "16: Research the best time of day to visit beaches\n",
      "17: Research the best time of day to take boat tours\n",
      "18: Research the best time of day to take parasailing tours\n",
      "19: Research the best time of day to take scuba diving tours\n",
      "20: Research the best time of day to take windsurfing tours\n",
      "21: Research the safety of outdoor\n",
      "\u001b[92m\u001b[1m\n",
      "*****NEXT TASK*****\n",
      "\u001b[0m\u001b[0m\n",
      "3: Create a map of Boracay\n",
      "\u001b[93m\u001b[1m\n",
      "*****TASK RESULT*****\n",
      "\u001b[0m\u001b[0m\n",
      "\n",
      "\n",
      "I will create a map of Boracay that includes all the outdoor activities and restaurants that serve Filipino food. The map will include the timeline and to-do list that you have already created, so that you can easily plan your trip. The map will also include the locations of all the beaches, parks, and other attractions in Boracay.\n",
      "\u001b[91m\u001b[1m\n",
      "*****TASK ENDING*****\n",
      "\u001b[0m\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'objective': 'Plan a travel itinerary to boracay for 5 days. Must include outdoor activities and eating filipino foods.'}"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "\n",
    "input_prompt = \"Human: \"\n",
    "query = input(input_prompt)\n",
    "\n",
    "text = f\"{input_prompt} {query}\"\n",
    "print(textwrap.fill(text, width=80))\n",
    "baby_agi({\"objective\": query})\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
